"""
更多用于 GUI 环境的辅助函数和对象

date: 2025/9/29
author: SiHeng Tang
file: lib_gui_client.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
from __future__ import annotations

import base64
import logging
from dataclasses import dataclass, field
from math import tan, radians
from typing import Any

import cv2
import numpy as np
from PySide6 import QtGui
from PySide6.QtCore import QAbstractItemModel, QModelIndex, QObject, Qt, Signal
from PySide6.QtCore import QMetaObject, QCoreApplication
from PySide6.QtWidgets import QGridLayout, QDialogButtonBox, QSizePolicy, QRadioButton, QComboBox

from helpers.geo_math_trans import GeoVector

logger = logging.getLogger(__name__)

START_UP = ("Hello, using this manual MJ RPC console with:\n"
            "    [.msg] for MJ RPC message method call.\n"
            "    [.mtd or white in default] for MJ RPC common method call.\n"
            "    to send list parameter with one element, please use SPACE in the parameter list.\n\n")

# 飞行模式，垂直起降变体和热气流模式不支持
AP_PLANE_MODE_NOT_SUPPORT = -1
AP_PLANE_MODE_MANUAL = 0
AP_PLANE_MODE_CIRCLE = 1
AP_PLANE_MODE_STABILIZE = 2
AP_PLANE_MODE_TRAINING = 3
AP_PLANE_MODE_ACRO = 4
AP_PLANE_MODE_FLY_BY_WIRE_A = 5
AP_PLANE_MODE_FLY_BY_WIRE_B = 6
AP_PLANE_MODE_CRUISE = 7
AP_PLANE_MODE_AUTOTUNE = 8
AP_PLANE_MODE_AUTO = 10
AP_PLANE_MODE_RTL = 11
AP_PLANE_MODE_LOITER = 12
AP_PLANE_MODE_TAKEOFF = 13
AP_PLANE_MODE_LAND = 14
AP_PLANE_MODE_GUIDED = 15
AP_PLANE_MODE_INITIALIZING = 16

# 飞行模式字符串，用于用户界面显示和日志
PLANE_FLIGHT_MODES_STRING = {
    AP_PLANE_MODE_NOT_SUPPORT: "NOT_SUPPORT",
    AP_PLANE_MODE_MANUAL: "MANUAL",
    AP_PLANE_MODE_CIRCLE: "CIRCLE",
    AP_PLANE_MODE_STABILIZE: "STABILIZE",
    AP_PLANE_MODE_TRAINING: "TRAINING",
    AP_PLANE_MODE_ACRO: "ACRO",
    AP_PLANE_MODE_FLY_BY_WIRE_A: "FLY_BY_WIRE_A",
    AP_PLANE_MODE_FLY_BY_WIRE_B: "FLY_BY_WIRE_B",
    AP_PLANE_MODE_CRUISE: "CRUISE",
    AP_PLANE_MODE_AUTOTUNE: "AUTOTUNE",
    AP_PLANE_MODE_AUTO: "AUTO",
    AP_PLANE_MODE_RTL: "RTL",
    AP_PLANE_MODE_LOITER: "LOITER",
    AP_PLANE_MODE_TAKEOFF: "TAKEOFF",
    AP_PLANE_MODE_LAND: "LAND",
    AP_PLANE_MODE_GUIDED: "GUIDED",
    AP_PLANE_MODE_INITIALIZING: "INITIALIZING",
}


def map_fm(_c: int):
    """映射飞行模式整数为一个单独的字母"""
    if _c in PLANE_FLIGHT_MODES_STRING:
        return PLANE_FLIGHT_MODES_STRING.get(_c)[0]
    else:
        return PLANE_FLIGHT_MODES_STRING[AP_PLANE_MODE_NOT_SUPPORT][0]


def cv2q_image(cv_img):
    """
    将 opencv 绘制图像转换到 Qt 图像用于显示
    :param cv_img: opencv 类矩阵输出
    :return: Qt 图像类
    """
    rgb_img = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)
    h, w, ch = rgb_img.shape
    bytes_per_line = ch * w
    return QtGui.QImage(rgb_img.data, w, h, bytes_per_line, QtGui.QImage.Format.Format_RGB888)


def string2frame(_s: str) -> np.ndarray:
    """
    将 base64 编码的 JPEG 图像转换为 opencv 图像
    :param _s: base64 字符串
    :return: opencv 图像，即 ndarray
    """
    raw_array = np.frombuffer(base64.b64decode(_s), np.uint8)
    cv_image = cv2.imdecode(raw_array, cv2.IMREAD_COLOR)

    if cv_image is None:
        raise RuntimeError("image decode failed")

    return cv_image


# auto generate by pyside6-uic not edit this manually
# noinspection PyPep8Naming,PyAttributeOutsideInit
class Ui_connections_dialog(object):

    def setupUi(self, connections_dialog):
        if not connections_dialog.objectName():
            connections_dialog.setObjectName(u"connections_dialog")
        connections_dialog.resize(350, 105)
        self.gridLayout = QGridLayout(connections_dialog)
        self.gridLayout.setObjectName(u"gridLayout")
        self.buttonBox = QDialogButtonBox(connections_dialog)
        self.buttonBox.setObjectName(u"buttonBox")
        sizePolicy = QSizePolicy(QSizePolicy.Policy.MinimumExpanding, QSizePolicy.Policy.Fixed)
        sizePolicy.setHorizontalStretch(0)
        sizePolicy.setVerticalStretch(0)
        sizePolicy.setHeightForWidth(self.buttonBox.sizePolicy().hasHeightForWidth())
        self.buttonBox.setSizePolicy(sizePolicy)
        self.buttonBox.setOrientation(Qt.Orientation.Horizontal)
        self.buttonBox.setStandardButtons(QDialogButtonBox.StandardButton.Cancel | QDialogButtonBox.StandardButton.Ok)

        self.gridLayout.addWidget(self.buttonBox, 2, 0, 1, 1)

        self.is_auto_reconnect = QRadioButton(connections_dialog)
        self.is_auto_reconnect.setObjectName(u"is_auto_reconnect")

        self.gridLayout.addWidget(self.is_auto_reconnect, 1, 0, 1, 1)

        self.com_list = QComboBox(connections_dialog)
        self.com_list.setObjectName(u"com_list")
        sizePolicy1 = QSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        sizePolicy1.setHorizontalStretch(0)
        sizePolicy1.setVerticalStretch(0)
        sizePolicy1.setHeightForWidth(self.com_list.sizePolicy().hasHeightForWidth())
        self.com_list.setSizePolicy(sizePolicy1)

        self.gridLayout.addWidget(self.com_list, 0, 0, 1, 1)

        self.gridLayout.setRowStretch(0, 1)
        self.gridLayout.setRowStretch(1, 1)

        self.retranslateUi(connections_dialog)
        self.buttonBox.accepted.connect(connections_dialog.accept)
        self.buttonBox.rejected.connect(connections_dialog.reject)

        QMetaObject.connectSlotsByName(connections_dialog)

    def retranslateUi(self, connections_dialog):
        connections_dialog.setWindowTitle(QCoreApplication.translate("connections_dialog", u"Connections", None))
        self.is_auto_reconnect.setText(
            QCoreApplication.translate("connections_dialog", u"Auto re-connect when connection error", None))


class TargetsSet(QObject):
    """
    目标集合，提供储存，增量更新和访问图片的功能
    """
    new_target = Signal()
    new_target_fast_key = Signal()

    def __init__(self):
        super().__init__()

        self._raw_targets: dict[int, tuple[Any, float, float, str]] = {}

    def add_target(self, pid: int, graphics: Any, lat_float: float, lon_float: float):
        """
        添加目标
        :param pid: 目标 ID，由飞行器端侧分配
        :param graphics: 图像，为 None 则不显示
        :param lat_float: latitude
        :param lon_float: longitude
        :return: 新目标返回 True，存在则返回 False
        """
        if pid in self._raw_targets:
            return False
        else:
            self._raw_targets.setdefault(pid, (graphics, lat_float, lon_float, ""))
            self.new_target.emit()
            return True

    def set_fast_key_seq_str(self, pid: int, seq_str: str):
        if pid in self._raw_targets:
            _g, _lat, _lon, _ = self._raw_targets.get(pid)
            self._raw_targets.pop(pid)
            self._raw_targets.setdefault(pid, (_g, _lat, _lon, seq_str))
            self.new_target_fast_key.emit()
            return True
        else:
            return False

    def get_target_obj_by_pid(self, pid: int):
        return self._raw_targets.get(pid, None)

    @property
    def target_pids(self):
        return [i for i in self._raw_targets.keys()]

    @property
    def positions_for_view(self):
        return [[k, v[1], v[2]] for k, v in self._raw_targets.items()]


@dataclass
class Node:
    key: str | int
    value: Any = ""
    value_type: type | None = None
    children: list[Node] = field(default_factory=list)


class JsonModel(QAbstractItemModel):
    data_modified = Signal(dict)

    def __init__(self, parent: QObject | None = None):
        super().__init__(parent)
        self._root = Node(key="root")
        self._headers = ("Key", "Value")

    def clear(self) -> None:
        self.load({})

    def load(self, data: dict) -> bool:
        """从字典加载JSON数据"""
        self.beginResetModel()
        self._root = self._build_node("root", data)
        self.endResetModel()
        self.data_modified.emit(data)
        return True

    def _build_node(self, key: str | int, value: Any) -> Node:
        """递归构建节点树"""
        node = Node(key=key, value_type=type(value))

        if isinstance(value, dict):
            node.children = [self._build_node(k, v) for k, v in value.items()]
        elif isinstance(value, list):
            node.children = [self._build_node(i, v) for i, v in enumerate(value)]
        else:
            node.value = value

        return node

    ##############################
    # QAbstractItemModel 接口实现
    ##############################
    def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole):
        if not index.isValid():
            return None

        node: Node = index.internalPointer()
        if role == Qt.ItemDataRole.DisplayRole:
            return node.key if index.column() == 0 else node.value
        if role == Qt.ItemDataRole.EditRole and index.column() == 1:
            return node.value
        return None

    def setData(self, index: QModelIndex, value: Any, role: int = Qt.ItemDataRole.EditRole):
        if role == Qt.ItemDataRole.EditRole and index.column() == 1:
            node: Node = index.internalPointer()
            node.value = value
            self.dataChanged.emit(index, index, [Qt.ItemDataRole.EditRole])

            logger.warning(f"Change config to {self.to_json()}")

            return True

        return False

    def headerData(self, section: int, orientation: Qt.Orientation, role: int = Qt.ItemDataRole.DisplayRole):
        if role == Qt.ItemDataRole.DisplayRole and orientation == Qt.Orientation.Horizontal:
            return self._headers[section]
        return None

    def index(self, row: int, column: int, parent: QModelIndex = QModelIndex()) -> QModelIndex:
        if not self.hasIndex(row, column, parent):
            return QModelIndex()

        parent_node = self._root if not parent.isValid() else parent.internalPointer()
        return self.createIndex(row, column, parent_node.children[row])

    def parent(self, index: QModelIndex) -> QModelIndex:
        if not index.isValid():
            return QModelIndex()

        child: Node = index.internalPointer()
        parent = self._find_parent(child, self._root)

        if parent is None or parent is self._root:
            return QModelIndex()

        return self.createIndex(parent.children.index(child), 0, parent)

    def _find_parent(self, target: Node, current: Node) -> Node | None:
        """查找目标节点的父节点"""
        for child in current.children:
            if child is target:
                return current
            found = self._find_parent(target, child)
            if found:
                return found
        return None

    def rowCount(self, parent: QModelIndex = QModelIndex()) -> int:
        if parent.column() > 0:
            return 0
        return len(self._root.children) if not parent.isValid() else len(parent.internalPointer().children)

    def columnCount(self, parent: QModelIndex = QModelIndex()) -> int:
        return 2

    def flags(self, index: QModelIndex) -> Qt.ItemFlags:
        flags = super().flags(index)
        return flags | Qt.ItemFlag.ItemIsEditable if index.column() == 1 else flags

    def to_json(self) -> dict:
        """将模型转换回JSON对象"""
        return self._node_to_json(self._root)

    def _node_to_json(self, node: Node) -> Any:
        if node.value_type is dict:
            return {child.key: self._node_to_json(child) for child in node.children}
        if node.value_type == list:
            return [self._node_to_json(child) for child in node.children]
        return node.value


class MapView(QObject):
    """
    地图拼图，将多张图像合成为一张图像
    """

    def __init__(self, fov_degrees=78, min_valid_alt=0.0, simplify_alt=1.0, max_memory_gb=0.5):
        """
        初始化拼图类
        :param fov_degrees: 相机视场角
        :param min_valid_alt: 最小确认高度
        :param simplify_alt: 最小图像处理高度，小于此高度的图像不处理
        :param max_memory_gb: 最大图像内存限制
        """
        super().__init__()
        self.max_memory_gb = max_memory_gb
        self.fov_rad = radians(fov_degrees)
        self.min_valid_alt = min_valid_alt
        self.simplify_alt = max(simplify_alt, 0.1)

        self.images = []
        self.min_x = None
        self.max_x = None
        self.min_y = None
        self.max_y = None
        self.center_geo = None
        self.target_positions = []

    @staticmethod
    def _rotate_image(img, heading_deg):
        """根据拍摄航向在地图上旋转图像"""
        heading_deg = 360 - heading_deg
        h, w = img.shape[:2]
        drag = int(np.ceil(np.sqrt(w * w + h * h)))
        big = np.zeros((drag, drag, 3), dtype=np.uint8)
        off_x = (drag - w) // 2
        off_y = (drag - h) // 2
        big[off_y:off_y + h, off_x:off_x + w] = img
        center = (drag // 2, drag // 2)
        m = cv2.getRotationMatrix2D(center, heading_deg, 1.0)
        rotated = cv2.warpAffine(big, m, (drag, drag),
                                 flags=cv2.INTER_LINEAR,
                                 borderValue=(0, 0, 0))
        return rotated, m, (drag, drag)

    @staticmethod
    def _get_rotated_corners_meters(rot_size, center_mx, center_my, gsd):
        """根据地面像素分辨率计算图像的四角的坐标"""
        rot_w, rot_h = rot_size
        half_w_m = (rot_w * gsd) / 2.0
        half_h_m = (rot_h * gsd) / 2.0
        return [
            (center_mx - half_w_m, center_my - half_h_m),
            (center_mx + half_w_m, center_my - half_h_m),
            (center_mx + half_w_m, center_my + half_h_m),
            (center_mx - half_w_m, center_my + half_h_m)
        ]

    def _memery_check(self):
        """检查内存占用，释放超过限制的帧，返回缓存的帧应有的拼图宽高，最精细的地面像素分辨率"""
        while True:
            min_gsd = min(item["gsd"] for item in self.images)
            width_px = int((self.max_x - self.min_x) / min_gsd) + 1
            height_px = int((self.max_y - self.min_y) / min_gsd) + 1
            memory_gb = (width_px * height_px * 3) / (1024 ** 3)

            if memory_gb <= self.max_memory_gb or len(self.images) <= 1:
                break

            self.images.pop(0)
            self._recalculate_bounds()
            logger.warning(
                f"Not enough memory for rendering ({memory_gb:.2f} GB needed, {self.max_memory_gb} GB allowed), removed oldest image")

        return width_px, height_px, min_gsd

    def _geo_to_meters(self, lat, lon):
        """计算当前图像与图像中心的距离"""
        if self.center_geo is None:
            self.center_geo = (lat, lon)
            return 0.0, 0.0

        center_gv = GeoVector(lat=self.center_geo[0], lon=self.center_geo[1])
        target_gv = GeoVector(lat=lat, lon=lon)
        diff_vector = center_gv.diff(target_gv)
        return diff_vector.dx, diff_vector.dy

    def _scale_image(self, img, alt):
        """计算地面像素分辨率"""
        if alt < self.simplify_alt:
            alt = self.simplify_alt

        h, w = img.shape[:2]
        ground_width = 2 * alt * tan(self.fov_rad / 2)
        gsd = ground_width / w
        return img.copy(), gsd

    def _recalculate_bounds(self):
        """重新计算拼图边界"""
        if not self.images:
            self.min_x = self.max_x = self.min_y = self.max_y = None
            return

        self.min_x = min(item["mx"] for item in self.images)
        self.max_x = max(item["mx"] for item in self.images)
        self.min_y = min(item["my"] for item in self.images)
        self.max_y = max(item["my"] for item in self.images)

        for item in self.images:
            half_w_m = (item["w"] * item["gsd"]) / 2.0
            half_h_m = (item["h"] * item["gsd"]) / 2.0
            self.min_x = min(self.min_x, item["mx"] - half_w_m)
            self.max_x = max(self.max_x, item["mx"] + half_w_m)
            self.min_y = min(self.min_y, item["my"] - half_h_m)
            self.max_y = max(self.max_y, item["my"] + half_h_m)

    def add_position_img(self, data):
        """
        将 scanner 的结果数据包添加到图像序列中
        :param data: scanner 结果字典对象
        """
        img = string2frame(data["overview"])
        pos = data["position"]
        lat, lon = pos["lat"], pos["lon"]
        heading = pos["hdg"]
        alt = pos["alt"]

        if alt < self.min_valid_alt:
            return

        scaled_img, gsd = self._scale_image(img, alt)
        rotated_img, m, rot_size = self._rotate_image(scaled_img, heading)

        mx, my = self._geo_to_meters(lat, lon)

        corners_m = self._get_rotated_corners_meters(rot_size, mx, my, gsd)
        for (x, y) in corners_m:
            self.min_x = x if self.min_x is None else min(self.min_x, x)
            self.max_x = x if self.max_x is None else max(self.max_x, x)
            self.min_y = y if self.min_y is None else min(self.min_y, y)
            self.max_y = y if self.max_y is None else max(self.max_y, y)

        self.images.append({
            "img": rotated_img,
            "mx": mx,
            "my": my,
            "w": rot_size[0],
            "h": rot_size[1],
            "gsd": gsd
        })

        self._memery_check()

        return img.copy()

    def get_img(self, max_image: int = -1):
        """
        返回拼图图像（含目标红叉及ID文本）
        :param max_image: 最大图像长度，设为负值选中整个图像
        :return: opencv 图像
        """
        if not self.images or self.min_x is None:
            return None

        width_px, height_px, min_gsd = self._memery_check()

        canvas = np.zeros((height_px, width_px, 3), dtype=np.uint8)

        idx_start = 0
        if max_image <= len(self.images) and idx_start > 0:
            idx_start = len(self.images) - max_image

        for item in self.images[idx_start:]:
            img = item["img"]
            mx, my = item["mx"], item["my"]
            h, w = img.shape[:2]

            px = int((mx - self.min_x) / min_gsd)
            py = int((self.max_y - my) / min_gsd)

            x1 = max(0, px - w // 2)
            y1 = max(0, py - h // 2)
            x2 = min(width_px, x1 + w)
            y2 = min(height_px, y1 + h)

            img_x1 = max(0, (px - w // 2) - x1)
            img_y1 = max(0, (py - h // 2) - y1)
            img_x2 = img_x1 + (x2 - x1)
            img_y2 = img_y1 + (y2 - y1)

            roi = canvas[y1:y2, x1:x2]
            img_crop = img[img_y1:img_y2, img_x1:img_x2]

            if img_crop.shape[:2] != roi.shape[:2]:
                img_crop = cv2.resize(img_crop, (roi.shape[1], roi.shape[0]), interpolation=cv2.INTER_LANCZOS4)

            mask = np.any(img_crop != 0, axis=-1).astype(np.uint8) * 255
            roi[:] = cv2.bitwise_and(roi, roi, mask=cv2.bitwise_not(mask))
            roi[:] += cv2.bitwise_and(img_crop, img_crop, mask=mask)

        if self.target_positions and self.center_geo is not None:
            base_size_px = int(1 / min_gsd)
            cross_size_px = max(base_size_px, 2)
            font_size = max(base_size_px / 8, 0.4)
            font_thickness = max(int(base_size_px / 10), 1)
            font = cv2.FONT_HERSHEY_SIMPLEX

            for [target_id, target_lat, target_lon] in self.target_positions:
                target_mx, target_my = self._geo_to_meters(target_lat, target_lon)
                if not (self.min_x <= target_mx <= self.max_x and self.min_y <= target_my <= self.max_y):
                    continue

                target_px = int((target_mx - self.min_x) / min_gsd)
                target_py = int((self.max_y - target_my) / min_gsd)

                cv2.line(canvas,
                         (target_px - cross_size_px // 2, target_py),
                         (target_px + cross_size_px // 2, target_py),
                         (0, 0, 255), 2)
                cv2.line(canvas,
                         (target_px, target_py - cross_size_px // 2),
                         (target_px, target_py + cross_size_px // 2),
                         (0, 0, 255), 2)

                id_text = str(int(target_id)) if isinstance(target_id, float) and target_id.is_integer() else str(
                    target_id)
                text_y = target_py - cross_size_px // 2 - 1
                text_width, _ = cv2.getTextSize(id_text, font, font_size, font_thickness)[0]
                text_x = target_px - text_width // 2
                stroke_offset = max(int(font_size), 1)
                for dx, dy in [(-stroke_offset, 0), (stroke_offset, 0), (0, -stroke_offset), (0, stroke_offset)]:
                    cv2.putText(canvas, id_text, (text_x + dx, text_y + dy),
                                font, font_size, (0, 0, 0), font_thickness, cv2.LINE_AA)
                cv2.putText(canvas, id_text, (text_x, text_y),
                            font, font_size, (255, 255, 255), font_thickness, cv2.LINE_AA)

        return canvas

    def set_target_positions(self, target_list):
        """
        设置目标列表，格式为 [[id, lat, lon], ...]
        :param target_list: 二维列表，id为数字（int/float），lat/lon为经纬度
        """
        valid_targets = []

        for target in target_list:
            # 校验格式：必须是 [id, lat, lon]，id为数字，lat/lon为数字
            if isinstance(target, (list, tuple)) and len(target) == 3:
                target_id, lat, lon = target
                if (isinstance(target_id, (int, float)) and
                        isinstance(lat, (int, float)) and
                        isinstance(lon, (int, float))):
                    valid_targets.append([target_id, lat, lon])

        self.target_positions = valid_targets
